package org.loong.crypto.extension.cms;

import java.io.IOException;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.cert.CRLException;
import java.security.cert.CertificateEncodingException;
import java.security.cert.X509Certificate;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.Iterator;
import java.util.Objects;

import org.apache.commons.lang3.ArrayUtils;
import org.bouncycastle.asn1.ASN1Encoding;
import org.bouncycastle.asn1.ASN1ObjectIdentifier;
import org.bouncycastle.asn1.cms.CMSObjectIdentifiers;
import org.bouncycastle.cert.X509CertificateHolder;
import org.bouncycastle.cert.jcajce.JcaCRLStore;
import org.bouncycastle.cert.jcajce.JcaCertStore;
import org.bouncycastle.cms.CMSException;
import org.bouncycastle.cms.CMSProcessableByteArray;
import org.bouncycastle.cms.CMSSignedData;
import org.bouncycastle.cms.CMSSignedDataGenerator;
import org.bouncycastle.cms.CMSTypedData;
import org.bouncycastle.cms.DefaultSignedAttributeTableGenerator;
import org.bouncycastle.cms.SignerInfoGenerator;
import org.bouncycastle.cms.SignerInformation;
import org.bouncycastle.cms.SignerInformationVerifier;
import org.bouncycastle.cms.jcajce.JcaSignerInfoGeneratorBuilder;
import org.bouncycastle.operator.BufferingContentSigner;
import org.bouncycastle.operator.ContentSigner;
import org.bouncycastle.operator.OperatorCreationException;
import org.loong.crypto.common.exception.CryptoException;
import org.loong.crypto.core.algorithm.SignatureAlgorithm;
import org.loong.crypto.core.utils.ECKeyUtils;
import org.loong.crypto.core.utils.KeyTools;
import org.loong.crypto.extension.cms.jce.JceSignerInfoVerifierBuilder;
import org.loong.crypto.extension.operator.jce.JceContentSignerBuilder;
import org.loong.crypto.extension.operator.jce.JceDigestCalculatorProviderBuilder;
import org.loong.crypto.service.CryptoService;
import org.loong.crypto.spi.Join;

/**
 * The signedData provider.
 */
@Join
public class SignedDataProvider implements CMSProvider {

    @Override
    public byte[] generate(CryptoService cryptoService, CMSParameters params) throws CryptoException {
        if (ArrayUtils.isEmpty(params.getData())) {
            throw new CryptoException("待签名的数据不允许为空.");
        }

        if (Objects.isNull(params.getSignatureAlgorithm())) {
            throw new CryptoException("数据签名算法不允许为空.");
        }

        final X509Certificate signingCertificate = params.getSigningCertificate();
        final PublicKey signingPublicKey = params.getSigningPublicKey();
        if (Objects.isNull(signingCertificate) && Objects.isNull(signingPublicKey)) {
            throw new CryptoException("签名者证书或公钥不允许为空.");
        }

        final PrivateKey privateKey = params.getPrivateKey();
        if (Objects.isNull(privateKey)) {
            throw new CryptoException("签名者私钥不允许为空.");
        }

        try {
            final SignatureAlgorithm signatureAlgorithm = SignatureAlgorithm.find(params.getSignatureAlgorithm());
            final JceContentSignerBuilder contentSignerBuilder = new JceContentSignerBuilder(cryptoService, signatureAlgorithm.getName());
            final ContentSigner signer = new BufferingContentSigner(contentSignerBuilder.build(privateKey), 20480);
            final JcaSignerInfoGeneratorBuilder signerInfoGeneratorBuilder = new JcaSignerInfoGeneratorBuilder(new JceDigestCalculatorProviderBuilder(cryptoService).build())
                    .setSignedAttributeGenerator(new DefaultSignedAttributeTableGenerator()).setDirectSignature(params.isHasNoSignedAttributes());
            SignerInfoGenerator signerInfoGenerator = null;
            if (Objects.nonNull(signingCertificate)) {
                signerInfoGenerator = signerInfoGeneratorBuilder.build(signer, signingCertificate);
            } else if (Objects.nonNull(signingPublicKey)) {
                byte[] keyId = KeyTools.createSubjectKeyId(signingPublicKey).getKeyIdentifier();
                signerInfoGenerator = signerInfoGeneratorBuilder.build(signer, keyId);
            }

            final CMSSignedDataGenerator signedDataGenerator = new CMSSignedDataGenerator();
            signedDataGenerator.addSignerInfoGenerator(signerInfoGenerator);

            if (Objects.nonNull(signingCertificate)) {
                signedDataGenerator.addCertificates(new JcaCertStore(Collections.singletonList(signingCertificate)));
            }

            if (ArrayUtils.isNotEmpty(params.getCrls())) {
                signedDataGenerator.addCRLs(new JcaCRLStore(Arrays.asList(params.getCrls())));
            }

            ASN1ObjectIdentifier contentType = CMSObjectIdentifiers.signedData;
            if (ECKeyUtils.isSM2Algorithm(signatureAlgorithm.getOID())) {
                contentType = GMCMSObjectIdentifiers.signedData;
            }

            CMSTypedData cmsData = new CMSProcessableByteArray(contentType, params.getData());
            CMSSignedData signedData = signedDataGenerator.generate(cmsData, true);
            return signedData.getEncoded(ASN1Encoding.DER);
        } catch (CertificateEncodingException | CMSException | NoSuchAlgorithmException | OperatorCreationException | IOException | CRLException e) {
            throw new CryptoException(e.getMessage(), e);
        }
    }

    @Override
    public byte[] parse(CryptoService cryptoService, CMSParameters params) throws CryptoException {
        if (ArrayUtils.isEmpty(params.getData())) {
            throw new CryptoException("待解密的数据不允许为空.");
        }

        final X509Certificate signingCertificate = params.getSigningCertificate();
        final PublicKey signingPublicKey = params.getSigningPublicKey();
        if (Objects.isNull(signingCertificate) && Objects.isNull(signingPublicKey)) {
            throw new CryptoException("签名者证书或公钥不允许为空.");
        }

        try {
            final CMSSignedData signedData = new CMSSignedData(params.getData());
            if (Objects.isNull(signedData)) {
                throw new CryptoException("签名数据内容为空.");
            }

            final CMSTypedData signedContent = signedData.getSignedContent();
            if (!CMSObjectIdentifiers.signedData.equals(signedContent.getContentType()) && !GMCMSObjectIdentifiers.signedData.equals(signedContent.getContentType())) {
                throw new CryptoException("签名数据类型错误.");
            }

            final JceSignerInfoVerifierBuilder signerInfoVerifierBuilder = new JceSignerInfoVerifierBuilder(new JceDigestCalculatorProviderBuilder(cryptoService).build());
            X509CertificateHolder signingCertificateHolder = null;
            byte[] keyId = null;
            SignerInformationVerifier signerInformationVerifier = null;
            if (Objects.nonNull(signingCertificate)) {
                signingCertificateHolder = new X509CertificateHolder(signingCertificate.getEncoded());
                keyId = KeyTools.createSubjectKeyId(signingCertificate.getPublicKey()).getKeyIdentifier();
                signerInformationVerifier = signerInfoVerifierBuilder.build(cryptoService, signingCertificate);
            } else if (Objects.nonNull(signingPublicKey)) {
                keyId = KeyTools.createSubjectKeyId(signingPublicKey).getKeyIdentifier();
                signerInformationVerifier = signerInfoVerifierBuilder.build(cryptoService, signingPublicKey);
            }

            boolean signValid = false;
            Collection<SignerInformation> signers = signedData.getSignerInfos().getSigners();
            if (signers.isEmpty()) {
                throw new CryptoException("签名数据签名者信息为空.");
            }

            Iterator<?> it = signers.iterator();
            while (it.hasNext() && Boolean.FALSE.equals(signValid)) {
                final SignerInformation signerInfo = (SignerInformation) it.next();
                if ((Objects.nonNull(signingCertificateHolder) && signerInfo.getSID().match(signingCertificateHolder))
                        || (Objects.nonNull(keyId) && Arrays.equals(signerInfo.getSID().getSubjectKeyIdentifier(), keyId))) {
                    signValid = signerInfo.verify(signerInformationVerifier);
                }
            }

            if (Boolean.FALSE.equals(signValid)) {
                throw new CryptoException("Signature validation failed.");
            }

            return (byte[]) signedContent.getContent();
        } catch (IOException | CertificateEncodingException | CMSException | OperatorCreationException e) {
            throw new CryptoException(e.getMessage(), e);
        }
    }
}
